#!/usr/bin/env python
#=======================================================================
#
#          FILE: fw-mlx-update.py
#         USAGE: ./fw-mlx-update.py
#   DESCRIPTION: Mellanox HCAs Automatic Firmware Update Tool
#  REQUIREMENTS: mstflint
#          BUGS: N/A
#         NOTES: MLNX_PUBLIC
#        AUTHOR: Ali Ayoub
#         EMAIL: ali@mellanox.com
#       COMPANY: Mellanox
#       CREATED: 09.17.2010-16:14:58
#      REVISION: 2.0.4000
#=======================================================================

import sys
import os
import time
#import popen2
import string

SCRIPT_VER = "2.0.4000"
SCRIPT_NAME = os.path.basename(sys.argv[0])
HELP = \
"Mellanox HCAs Automatic Firmware Update Tool v%s\n" % SCRIPT_VER + \
" this tool downloads FW image(s) for installed device(s)\n" + \
" from Mellanox website, and updates the FW (using mstflint)\n"

VERBOSE = 0
NOBURN = 0
OUTPUT_RED = " &> /dev/null"
ARCH = os.popen('uname -i').read().strip()
PROXY = None
MLNX_COM = "http://mellanox.com"
OPENFABRICS_COM = "http://www.openfabrics.org"
MST_BIN = "mstflint"
MST_FLAGS = None
MFT_WEB_SOURCE = MLNX_COM + "/downloads/MFT"
MST_WEB_SOURCE = OPENFABRICS_COM + "/downloads/mstflint"
# e.g.: http://www.openfabrics.org/downloads/mstflint/mstflint-1.4.tar.gz
OFED_WEB_SOURCE = OPENFABRICS_COM + "/downloads/OFED"
WGET_BIN = "wget"
WGET_WEB_SOURCE = "http://ftp.gnu.org/pub/gnu/wget"
UTILS = [MST_BIN, WGET_BIN, 'cat', 'grep', 'lspci', 'awk', 'unzip', 'which', 'du']
DEVICE = None
FW_IMG = None
PSID = None
TMP_DIR = "/tmp"
WGET_LOG = TMP_DIR + os.sep + "wget.log"
INTER = "yes"
USAGE = "Usage: " + SCRIPT_NAME + " [--help|-h] [--debug|-d] [--query|-q]" + \
	"\n\t [--mst_bin <mstflint path>]" + \
	"\n\t [--mst_flags <flags> (default %s)]" % str(MST_FLAGS) + \
	"\n\t [--wget_bin <wget path>]" + \
	"\n\t [--device <pciid to burn>]" + \
	"\n\t [--proxy <URL:PORT> (default %s)]" % str(PROXY) + \
	"\n\t [--fw_img <FW binary>]" + \
	"\n\t [--psid <PSID>]" + \
	"\n\t [--interactive <yes|no> (default %s)]" % INTER + \
	"\n"
##############################################################################
def vprint(msg):
        if VERBOSE > 0:
		print ""
                print msg

def remove(fn):
	vprint("-V- removing %s" % fn)
	try:
		os.remove(fn)
	except:
		pass

def get_pcids_desc():
	cmd = "lspci | grep -i mellanox | grep -v bridge"
	vprint("-V- cmd: " + cmd)
	pcids = os.popen(cmd).read().strip().split("\n")
	if pcids == "":
		return None
	return pcids

def get_pcids():
	pcids = []
	pcids_desc = get_pcids_desc()
	if pcids_desc == None:		
		return None
	vprint("-V- pcids_desc = " + str(pcids_desc))
	for p in pcids_desc:
		if not p.split(" ")[0] == '':
			pcids += [p.split(" ")[0]]
	return pcids

def get_name(pcid):
	cmd = "%s -d %s dc |  grep Description" % (MST_BIN, pcid)
	vprint("-V- cmd: " + cmd)
	out = os.popen(cmd).read().strip()
	return out.split('=')[1].strip()


def get_psid(pcid):
	cmd = MST_BIN + " -d " + pcid + " q | grep -i 'board id' "
	vprint("-V- cmd: " + cmd)
	psid = os.popen(cmd).read().strip().replace("(","").replace(")","").strip()
	if psid == "":
		return None
	psid = psid[psid.find("MT_"):]
	if len(psid) != 13:
		vprint("-V- psid too short (%s)" % psid)
		return None
	return psid


def get_page(url):
	cmd = WGET_BIN + " --timeout 8 -t 2 " + url + \
	' -o ' + WGET_LOG + ' && grep saved ' + WGET_LOG + \
	' | cut -d")" -f2 | cut -d- -f2- | cut -d" " -f2'

	vprint("-V- cmd: " + cmd)
	page = os.popen(cmd).read().strip()
	remove(WGET_LOG)
	if page == "":
		return None
	_page = os.path.basename(url)
	vprint("-V- check _page %s" % _page)
	if os.path.exists(_page):
		return _page
	_page = page
	vprint("-V- check _page %s" % _page)
	if os.path.exists(_page):
		return _page
	vprint("-V- get_page failed")
	return None

def get_href(line):
	return line.split('href="')[1].split('">')[0].strip()

def get_bin_url(psid):
	url_base = MLNX_COM + "/content"
	url_file = "pages.php?pg=firmware_download"
	url = url_base + os.sep + url_file
	tables_page = get_page(url)
	if tables_page == None:
		return tables_page
	cmd = "grep firmware_table %s" % url_file
	tables_lines = os.popen(cmd).readlines()
	remove(url_file)
	tables_urls = []
	for line in tables_lines:
		if line.count("href") == 0:
			break
		table_url = get_href(line)
		tables_urls += [table_url]
		vprint("-V- table_url: " + table_url)

	fw_url = None
	tables_urls
	for tables_url in tables_urls:
		url = url_base + os.sep + tables_url
		fw_page = get_page(url)
		if fw_page == None:
			return fw_page
		vprint("-V- downloaded tables_url: " + url)
		vprint("-V- fw_page %s" % fw_page)
		cmd = "grep %s %s -A7" % (psid, fw_page)
		vprint("-V- cmd: " + cmd)
		fw_url_line = ''
		fw_url_lines = os.popen(cmd).readlines()
		remove(fw_page)
		fw_url_lines.reverse()
		for fw_url_line in fw_url_lines:
			fw_url_line = fw_url_line.strip()
			if fw_url_line.count('.zip'):
				vprint("-V- found fw_url_line" + fw_url_line)
				break
		if len(fw_url_line.strip()):
			fw_url = get_href(fw_url_line)
			break

	if fw_url == None:
		vprint("-V- could not find fw_url")
		return fw_url
	return url_base + os.sep + fw_url

def get_bin_direct(url):
	if not PROXY == None:
		os.environ['http_proxy']=str(PROXY)
	bin = get_page(url)
	if bin == None:
		return bin
	if bin.count(".zip"):
		cmd = "unzip -o " + bin
		vprint("-V- cmd: " + cmd)
		unzip_out = os.popen(cmd).read().strip()
		vprint("-V- unzip_out: " + unzip_out)
		bin_ext = unzip_out.split()[-1]
		os.system("rm -rf " + bin)
		bin = bin_ext
	return bin

def burn_bin(pcid, bin):
	cmd  = MST_BIN
	if MST_FLAGS != None:
		cmd += " " + MST_FLAGS + " "
	if INTER == "no":
		cmd += " -y"
	cmd += " -d " + pcid + " -i " + bin
	cmd += " b"
	vprint("-V- cmd: " + cmd)
	return os.system(cmd)

def check_rc(func, rc, err_msg = ""):
	if rc == None:
		print "-E- %s failed" % func
		if not err_msg == "":
			print "-I- %s" % err_msg
		sys.exit(1)

def parse_args (args):
	global VERBOSE
	global OUTPUT_RED
	global MST_BIN
	global MST_FLAGS
	global WGET_BIN
	global PROXY
	global DEVICE
	global NOBURN
	global FW_IMG
	global PSID
	global INTER

	index  = 0
	while index < len(args):
		arg = args[index]
		index = index +1
		if len(args) > index:
			val = args[index]
		else:
			val = None

		# loop on args
		if arg == "--debug" or arg == "-d":
			VERBOSE = 1
			OUTPUT_RED = ""
		elif arg == "--help" or arg == "-h":
			print HELP
			print USAGE			
			return 2
		elif arg == "--mst_bin":
			MST_BIN = val
			index = index +1
		elif arg == "--mst_flags":
			MST_FLAGS = val
			index = index +1
		elif arg == "--wget_bin":
			WGET_BIN = val
			index = index +1
		elif arg == "--proxy":			
			PROXY = val
			index = index +1
		elif arg == "--device":
			DEVICE = val
			index = index +1
		elif arg == "--query" or arg == "-q":
			NOBURN = 1
		elif arg == "--fw_img":
			FW_IMG = val
			index = index +1
		elif arg == "--psid":
			PSID = val
			index = index +1
		elif arg == "--interactive" and val in ['yes', 'no']:
			INTER = val
			index = index +1
		else:
			print "-E- Unrecognized parameter (%s)" % arg
			return 1
	return 0

def short_url(url):
	max = 45
	if len(url) <= max:
		return url
	return url[0:27] + ".." + url[-15:]

def getsize(fn):
	# I did't use os.path.getize() because it's output is not 'human' readable
	cmd = "\du -sh %s | awk '{print $1}'" % fn
	return os.popen(cmd).read().strip()

def checkfile(fn):
	cmd = "\which " + fn
	pfn = os.popen(cmd).read().strip()
	return os.path.exists(pfn)

def check_utils():
	for util in UTILS:
		if checkfile(util):
			continue
		rc = "%s not found" % util
		util_check = 'wget'
		if util.count(util_check):
			rc += " [Try --" + util_check +"_bin flag OR download " + util_check + " from %s]" % WGET_WEB_SOURCE
		util_check = 'mstflint'
		if util.count(util_check):
			rc += " [Try --" + util_check + "_bin flag OR download " + util_check + " from %s]" % MST_WEB_SOURCE
		return rc
	return None

def check_net():
	cmd = "cd %s; %s --timeout 10 %s -o %s" % (TMP_DIR, WGET_BIN, MLNX_COM, WGET_LOG) 
	vprint("-V- cmd: " + cmd)
	rc = os.system(cmd)
	if rc:
		cmd = "cat %s" % WGET_LOG
		vprint("-V- cmd: " + cmd)
		return os.popen(cmd).read().strip()
	remove(TMP_DIR + os.sep + WGET_LOG)
	remove(TMP_DIR + os.sep + "index.html")
	return None

def check_params():
	msg = None
	if DEVICE == None:
		if FW_IMG != None or PSID != None:
			msg = "Must specify DEVICE"
	if NOBURN:
		if PSID != None:
			msg = "Cannot query PSID"
	return msg

##############################################################################
if __name__ == '__main__':
	# check user
	if os.geteuid() != 0:
		print "-E- you must be root to run this script."
		sys.exit(1)

	# parse args, 0:ok, 1:err, 2:exit
	rc = parse_args(sys.argv[1:])
	if rc == 1:
		print "-E- Invalid input"
		sys.exit(rc)
	elif rc == 2:
		sys.exit(0)

	# check local utils
	rc = check_utils()
	if not rc == None:
		print "-E- Missing utility: %s" % rc
		sys.exit(1)

	# check internet
	rc = check_net()
	if not rc == None:
		print "-E- Internect connection problem: %s" % rc
		sys.exit(1)

	# check params
	rc = check_params()
	if not rc == None:
		print "-E- Bad parameter: %s" % rc
		sys.exit(1)

	# burn
	pcids = get_pcids()
	if pcids == None or pcids == []:
		print "-I- No HCAs present"
		sys.exit(0)

	check_rc("get_pcids", pcids)	
	vprint("-V- pcids = " + str(pcids))
	USER_PSID = PSID
	for pcid in pcids:
		if DEVICE != None and DEVICE != pcid:
			print "-I- skip %s (looking for %s)" % (pcid, DEVICE)
			continue
		print "Device PCI ID    " + str(pcid)

		PSID = USER_PSID
		if USER_PSID == None:
			PSID = get_psid(pcid)
			check_rc("get_psid", PSID)
		print "Device PSID      " + str(PSID)

		dev_name = get_name(pcid)
		print "Device Name      " + str(dev_name)

		if NOBURN:
			os.system("%s -d %s q" % (MST_BIN, pcid))
			print "-I- skip burning"
			print ""
			continue;

		vprint("-V- FW_IMG " + str(FW_IMG))
		if FW_IMG == None:
			sys.stdout.write("Searching URL... ")
			sys.stdout.flush()

			bin_url = get_bin_url(PSID)
			check_rc("get_bin_url", bin_url)
			string = short_url(bin_url)
			print str(string)

			if not PROXY == None:
				print "http_proxy       " + str(PROXY)

			sys.stdout.write("Downloading...   ")
			sys.stdout.flush()
		
			bin = get_bin_direct(bin_url)
			check_rc("get_bin_direct", bin, "try --proxy flag")
			string = bin + " (" + getsize(bin) + ")"
			print str(string)
		else:
			sys.stdout.write("Use FW image:    ")
			sys.stdout.flush()
		
			bin = str(FW_IMG)
			print str(bin)			

		sys.stdout.write("Burning...   ")
		sys.stdout.flush()
		burn_bin(pcid, bin) # ignore RC to allow burning other pcid's
		remove(bin)
		print ""

	print "Done"
	sys.exit(rc)
